# -*- coding: utf-8 -*-
"""
Author: 王薪（Wang, Xin）
E-mail: xin.wang.418@aliyun.com
August 2020
"""

from odbAccess import *
import os
import numpy as np
import re
from operator import attrgetter


class tocsv():

    def __init__(self, odbName, cwd, step=1):

        self.odbName = odbName
        self.cwd = cwd
        self.outd = os.path.join(cwd, 'results')
        self.numNodes = 0
        self.numElements = 0
        self.step = step

        if not os.path.exists(self.outd):
            os.makedirs(self.outd)

        self.odb = openOdb(path=os.path.join(self.cwd, self.odbName + '.odb'))

    def eleset(self):

        myAssembly = self.odb.rootAssembly
        numNodes = numElements = 0
        for name, instance in myAssembly.instances.items():
            numNodes = len(instance.nodes)
            numElements = len(instance.elements)

        instanceName = myAssembly.instances.keys()
        allEleSetName = myAssembly.instances[instanceName[0]].elementSets.keys()

        tempEleset = allEleSetName[0]
        temp = re.compile("([a-zA-Z]+)([0-9]+)")
        keyword = temp.match(tempEleset).groups()[0]

        eleSetName = []
        for word in allEleSetName:
            if word.startswith(keyword):
                eleSetName.append(word)

        # removeKeywords = ['GRAIN100_1', 'GRAIN110_1', 'GRAIN111_1']
        removeKeywords = []
        for word in removeKeywords:
            eleSetName.remove(word)

        eleSetNum = len(eleSetName)
        eleSetValue = np.zeros((numElements, 1))
        elementLabelArray = np.empty((0, 1))
        for i in range(0, eleSetNum):
            eleSet = myAssembly.instances[instanceName[0]].elementSets[eleSetName[i]]
            ele = eleSet.elements
            for eleNum in range(0, len(ele)):
                elementLabelArray = np.append(elementLabelArray, ele[eleNum].label)
            elementLabelArray = elementLabelArray.astype('int64')
            eleSetValue[elementLabelArray - 1] = i + 1
            elementLabelArray = np.empty((0, 1))

        np.savetxt(os.path.join(self.outd, '%s_MESET.csv' % self.odbName), eleSetValue, fmt='%i', delimiter='\t')
        print 'ELSET info has been saved.'

    def nodeElement(self):

        myAssembly = self.odb.rootAssembly

        stepNum = len(self.odb.steps.values())
        stepInfoArray = np.empty((0, 2))
        for i in range(stepNum):
            tempStepName = self.odb.steps.keys()[i]
            frameNum = len(self.odb.steps[tempStepName].frames)
            tempStepInfo = np.array([i + 1, frameNum])
            stepInfoArray = np.append(stepInfoArray, tempStepInfo)
        stepInfoArray = stepInfoArray.reshape((-1, 2))
        np.savetxt(os.path.join(self.outd, '%s_STEP.csv' % self.odbName), stepInfoArray, fmt='%i', delimiter='\t')


        for name, instance in myAssembly.instances.items():
            if instance.embeddedSpace == THREE_D:
                nodeDimension = 3
            else:
                nodeDimension = 2

        self.numNodes = len(instance.nodes)
        self.numElements = len(instance.elements)

        nodeArray = np.empty((0, nodeDimension))
        for node in instance.nodes:
            nodeArray = np.append(nodeArray, [node.coordinates], axis=0)
        np.savetxt(os.path.join(self.outd, '%s_MN.csv' % self.odbName), nodeArray, fmt='%.5f', delimiter='\t')
        print 'NODE info has been saved.'


        elementArray = np.empty((0, 1))
        for element in instance.elements:
            for nodeNum in element.connectivity:
                elementArray = np.append(elementArray, [nodeNum])
        elementArray = np.reshape(elementArray, (self.numElements, -1))
        np.savetxt(os.path.join(self.outd, '%s_ME.csv' % self.odbName), elementArray, fmt='%i', delimiter='\t')
        print 'ELEMENT info has been saved.'

        return self.numNodes, self.numElements

    def nodeElementDeformed(self):

        myAssembly = self.odb.rootAssembly

        numNodes = 0

        for name, instance in myAssembly.instances.items():
            numNodes = len(instance.nodes)
            nodeArray = np.empty((0, 3))
            for node in instance.nodes:
                nodeArray = np.append(nodeArray, [node.coordinates], axis=0)

        numStep = len(self.odb.steps)

        dataArray = np.zeros((1, 3))
        displacementArray = np.empty((0, 3))
        for i in range(0, numStep):
            numFrame = len(self.odb.steps[self.odb.steps.values()[i].name].frames)
            print '-----------------------------------------------------'
            print 'STEP%s has %s FRAMES' % (i + 1, numFrame)
            for j in lrange(0, numFrame - 1, self.step):
                frame = self.odb.steps[self.odb.steps.values()[i].name].frames[j]
                fieldOutputValues = frame.fieldOutputs['U'].values
                for d in fieldOutputValues:
                    dataArray = np.array(d.data).reshape(1, 3)
                    displacementArray = np.append(displacementArray, dataArray)
                displacementArray = displacementArray.reshape(self.numNodes, -1)
                displacementArray = displacementArray + nodeArray
                np.savetxt(os.path.join(self.outd, '%s_MN_s%s_f%s.csv' % (self.odbName, i + 1, j + 1)),
                           displacementArray,
                           fmt='%.5f', delimiter='\t')
                dataArray = np.zeros((1, 3))
                displacementArray = np.empty((0, 1))
                print 'NODE info in DEFORMED coordinates from STEP%s FRAME%s has been saved' % (i + 1, j + 1)
                j += 1
            i += 1
        print 'All NODE info in DEFORMED coordinates has been saved'

    def fieldOutput(self):

        myAssembly = self.odb.rootAssembly
        numStep = len(self.odb.steps)
        # -----------------------------------------------------

        # -----------------------------------------------------
        firstStepName = self.odb.steps.keys()[0]
        firstFrame = self.odb.steps['%s' % firstStepName].frames[0]

        if not os.path.isfile(os.path.join(self.outd, "fieldOutput.txt")):
            f = open(os.path.join(self.outd, "fieldOutput.txt"), "w")
            f.close()

        print 'Field outputs are listed below:'
        for fieldName in firstFrame.fieldOutputs.keys():
            print fieldName
        print '------------------------------------------------------------'
        inputFlag = int(raw_input('Which way do you want to specify field output variable?\n'
                                  '1. Read from file; \n'
                                  '2. Keyboard input;\n'
                                  '3. All the field outputs.\n'
                                  'Please input 1, 2 or 3: '))
        if inputFlag == 1:
            f = open(os.path.join(self.outd, "fieldOutput.txt"), "r")
            variable = f.readlines()[0]
            variable = [v.strip() for v in variable.split(',')]
        elif inputFlag == 2:
            variable = raw_input('Please select field output variables: ')
            variable = [v.strip() for v in variable.split(',')]
        elif inputFlag == 3:
            variable = firstFrame.fieldOutputs.keys()

        # -----------------------------------------------------
        for var in variable:
            for i in range(0, numStep):
                numFrame = len(self.odb.steps[self.odb.steps.values()[i].name].frames)
                print '------------------------------------------------------------'
                print 'STEP%s has %s FRAMES in total' % (i + 1, numFrame)
                for j in lrange(0, numFrame - 1, self.step):
                    frame = self.odb.steps[self.odb.steps.values()[i].name].frames[j]
                    fieldOutputObject = frame.fieldOutputs['%s' % var].bulkDataBlocks
                    fieldOutputArray = map(attrgetter('data'), fieldOutputObject)[0]
                    numComponents = fieldOutputArray.shape[1]
                    np.savetxt(os.path.join(self.outd, '%s_%s_s%s_f%s.csv' % (self.odbName, var, i + 1, j + 1)),
                               fieldOutputArray, fmt='\t'.join(['%.5f'] * numComponents))
                    print 'FIELD OUTPUT %s from STEP%s FRAME%s has been saved' % (var, i + 1, j + 1)
                    j += 1
                i += 1
        print 'All specified FIELD OUTPUT info has been saved'

    def fieldOutputUN(self, avgThreshold=0.75):

        myAssembly = self.odb.rootAssembly
        numStep = len(self.odb.steps)
        # -----------------------------------------------------

        # -----------------------------------------------------
        firstStepName = self.odb.steps.keys()[0]
        firstFrame = self.odb.steps['%s' % firstStepName].frames[0]

        if not os.path.isfile(os.path.join(self.outd, "fieldOutput.txt")):
            f = open(os.path.join(self.outd, "fieldOutput.txt"), "w")
            f.close()

        print 'Field outputs are listed below:'
        for fieldName in firstFrame.fieldOutputs.keys():
            print fieldName
        print '------------------------------------------------------------'
        inputFlag = int(raw_input('Which way do you want to specify field output variable?\n'
                                  '1. Read from file; \n'
                                  '2. Keyboard input;\n'
                                  '3. All the field outputs.\n'
                                  'Please input 1, 2 or 3: '))
        if inputFlag == 1:
            f = open(os.path.join(self.outd, "fieldOutput.txt"), "r")
            variable = f.readlines()[0]
            variable = [v.strip() for v in variable.split(',')]
        elif inputFlag == 2:
            variable = raw_input('Please select field output variables: ')
            variable = [v.strip() for v in variable.split(',')]
        elif inputFlag == 3:
            variable = firstFrame.fieldOutputs.keys()

        # -----------------------------------------------------
        for var in variable:
            for i in range(0, numStep):
                numFrame = len(self.odb.steps[self.odb.steps.values()[i].name].frames)
                print '------------------------------------------------------------'
                print 'STEP%s has %s FRAMES in total' % (i + 1, numFrame)
                for j in lrange(0, numFrame - 1, self.step):
                    frame = self.odb.steps[self.odb.steps.values()[i].name].frames[j]
                    fieldOutputObject = frame.fieldOutputs['%s' % var].getSubset(position=ELEMENT_NODAL)
                    fieldOutputObject = fieldOutputObject.bulkDataBlocks
                    inFOArray = map(attrgetter('data'), fieldOutputObject)[0]
                    nodeLablesArray = map(attrgetter('nodeLabels'), fieldOutputObject)[0]

                    nodeLabelUnq = np.unique(nodeLablesArray)
                    numNodes = nodeLabelUnq.size

                    if len(inFOArray.shape) == 1:
                        numComponents = 1
                    else:
                        numComponents = inFOArray.shape[1]

                    if numComponents == 1:

                        inFOMax = np.max(inFOArray)
                        inFOMin = np.min(inFOArray)

                        fieldOutputArray = np.zeros(numNodes)

                        if inFOMax - inFOMin == 0:
                            pass

                        else:
                            for ii in range(numNodes):
                                testBoolean = nodeLablesArray == ii + 1
                                tempNodeList = np.where(testBoolean)[0]

                                if len(tempNodeList) == 1:
                                    fieldOutputArray[ii] = np.take(inFOArray, tempNodeList)
                                else:
                                    tempFieldOutputArray = np.take(inFOArray, tempNodeList)
                                    tempMax = np.max(tempFieldOutputArray)
                                    tempMin = np.min(tempFieldOutputArray)
                                    if (tempMax - tempMin) / (inFOMax - inFOMin) < avgThreshold:
                                        fieldOutputArray[ii] = np.mean(tempFieldOutputArray)
                                    else:
                                        fieldOutputArray[ii] = np.max(tempFieldOutputArray)
                                ii += 1

                    else:
                        inFOMax = np.amax(inFOArray, axis=0)
                        inFOMin = np.amin(inFOArray, axis=0)

                        fieldOutputArray = np.zeros([numNodes, numComponents])

                        if 0 in (inFOMax - inFOMin):
                            pass
                        else:
                            for ii in range(numNodes):
                                testBoolean = nodeLablesArray == ii + 1
                                tempNodeList = np.where(testBoolean)[0]

                                if len(tempNodeList) == 1:
                                    fieldOutputArray[i] = np.take(inFOArray, tempNodeList, axis=0)
                                    ii += 1
                                else:
                                    tempFieldOutputArray = np.take(inFOArray, tempNodeList, axis=0)
                                    tempArray = np.zeros([1, numComponents])
                                    tempMax = np.amax(tempFieldOutputArray, axis=0)
                                    tempMin = np.amin(tempFieldOutputArray, axis=0)

                                    for ij in range(numComponents):
                                        if (tempMax[ij] - tempMin[ij]) / (
                                                inFOMax[ij] - inFOMin[ij]) < avgThreshold:
                                            tempArray[:, ij] = np.mean(tempFieldOutputArray[:, ij])
                                        else:
                                            tempArray[:, ij] = np.max(tempFieldOutputArray[:, ij])
                                        ij += 1
                                    fieldOutputArray[ii] = tempArray
                                ii += 1

                    np.savetxt(os.path.join(self.outd, '%s_%s_UN_s%s_f%s.csv' % (self.odbName, var, i + 1, j + 1)),
                               fieldOutputArray, fmt='\t'.join(['%.5f'] * numComponents))
                    print 'FIELD OUTPUT (UNIQUE NODAL) %s value from STEP%s FRAME%s has been saved' % (
                    var, i + 1, j + 1)
                    j += 1
                i += 1
        print 'All specified FIELD OUTPUT (UNIQUE NODAL) value have been saved'

    def getInvariants(self):

        print 'Save invariants function only works for STRESS and STRAIN type tensor'
        print 'Save order: maxPrincipal, midPrincipal, minPrincipal, mises , tresca, inv3, press'

        variable = raw_input('Please input field output variables:')
        variable = [v.strip() for v in variable.split(',')]

        numStep = len(self.odb.steps)

        for var in variable:
            for i in range(0, numStep):
                numFrame = len(self.odb.steps[self.odb.steps.values()[i].name].frames)
                print '------------------------------------------------------------'
                print 'STEP%s has %s FRAMES in total' % (i + 1, numFrame)
                for j in lrange(0, numFrame - 1, self.step):
                    frame = self.odb.steps[self.odb.steps.values()[i].name].frames[j]
                    fieldOutputObject = frame.fieldOutputs['%s' % var].values
                    maxPrincipalArray = np.array(map(attrgetter('maxPrincipal'), fieldOutputObject)).reshape(-1, 1)
                    midPrincipalArray = np.array(map(attrgetter('midPrincipal'), fieldOutputObject)).reshape(-1, 1)
                    minPrincipalArray = np.array(map(attrgetter('minPrincipal'), fieldOutputObject)).reshape(-1, 1)
                    misesArray = np.array(map(attrgetter('mises'), fieldOutputObject)).reshape(-1, 1)
                    trescaArray = np.array(map(attrgetter('tresca'), fieldOutputObject)).reshape(-1, 1)
                    inv3Array = np.array(map(attrgetter('inv3'), fieldOutputObject)).reshape(-1, 1)
                    pressArray = np.array(map(attrgetter('press'), fieldOutputObject)).reshape(-1, 1)
                    fieldOutputArray = np.concatenate((maxPrincipalArray, midPrincipalArray, minPrincipalArray,
                                                       misesArray, trescaArray, inv3Array, pressArray), axis=1)
                    numComponents = fieldOutputArray.shape[1]
                    np.savetxt(os.path.join(self.outd, '%s_%s_inv_s%s_f%s.csv' % (self.odbName, var, i + 1, j + 1)),
                               fieldOutputArray, fmt='\t'.join(['%.5f'] * numComponents))
                    print 'Invariants of %s from STEP%s FRAME%s has been saved' % (var, i + 1, j + 1)
                    j += 1
                i += 1
        print 'All specified INVARIANTS values have been saved'

    def getTime(self):
        numStep = len(self.odb.steps)
        frameInfo = np.empty([0, 2])
        for i in range(0, numStep):
            numFrame = len(self.odb.steps[self.odb.steps.values()[i].name].frames)
            for j in lrange(0, numFrame - 1, self.step):
                frame = self.odb.steps[self.odb.steps.values()[i].name].frames[j]
                incrementNumber = frame.incrementNumber
                stepTime = frame.frameValue
                print np.array([stepTime, incrementNumber])
                frameInfo = np.vstack((frameInfo, np.array([stepTime, incrementNumber])))
                j += 1

        np.savetxt(os.path.join(self.outd, '%s_FRAME.csv' % self.odbName), frameInfo, fmt='%i %.2f', delimiter='\t')
        print 'FRAME info has been saved'

    @property
    def historyOutput(self):
        numStep = len(self.odb.steps)
        for i in range(0, numStep):
            stepObject = self.odb.steps[self.odb.steps.values()[i].name]
            historyOutputRegion = stepObject.historyRegions.keys()
            historyOutputVars = stepObject.historyRegions[historyOutputRegion[0]].historyOutputs.keys()
            numTimeStep = np.array(stepObject.historyRegions[historyOutputRegion[0]].historyOutputs[historyOutputVars[0]].data).shape[0]
            for hVar in historyOutputVars:
                historyOutputArray = np.empty([numTimeStep, 0])
                for rVar in historyOutputRegion:
                    historyOutputRepo = stepObject.historyRegions[rVar].historyOutputs
                    historyOutputObject = historyOutputRepo[hVar]
                    tempArray = np.array(historyOutputObject.data)
                    historyOutputArray = np.concatenate((historyOutputArray, tempArray[:, 1].reshape(-1, 1)), axis=1)
                if 'U' in hVar:
                    historyOutputArray = np.mean(historyOutputArray, axis=1).reshape(-1, 1)
                else:
                    historyOutputArray = np.sum(historyOutputArray, axis=1).reshape(-1, 1)

                historyOutputArray = np.concatenate((tempArray[:, 0].reshape(-1, 1), historyOutputArray), axis=1)

                np.savetxt(os.path.join(self.outd, '%s_HO_%s.csv' % (self.odbName, hVar)), historyOutputArray,
                           fmt='%.2f %.5f', delimiter='\t')

                print 'HISTORY OUTPUT %s value has been saved' % hVar

        historyOutputConArray = np.empty([numTimeStep, 0])
        numComponents = len(historyOutputVars)
        for hVar in historyOutputVars:
            tempArray = np.loadtxt(os.path.join(self.outd, '%s_HO_%s.csv' % (self.odbName, hVar)))
            historyOutputConArray = np.concatenate((historyOutputConArray, tempArray[:, 1].reshape(-1, 1)), axis=1)
        historyOutputConArray = np.concatenate((tempArray[:, 0].reshape(-1, 1), historyOutputConArray), axis=1)
        np.savetxt(os.path.join(self.outd, '%s_HO_CON.csv' % self.odbName), historyOutputConArray, fmt='\t'.join(['%.2f'] + ['%.5f']*numComponents))
        print 'Concatenated HISTORY OUTPUT values have been saved!'


def lrange(start, end, step):
    i = start
    while i < end:
        yield i
        i += step
    yield end
